# srt_textgrid_app.py
# 使用 Gradio 构建 Web 界面，支持 SRT ↔ TextGrid 双向转换

import gradio as gr
import os
import re

# ========== 复用之前的转换函数 ==========

def srt_time_to_seconds(srt_time: str) -> float:
    time_str = srt_time.replace(',', ':')
    parts = time_str.split(':')
    hours, minutes, seconds, milliseconds = map(float, parts)
    return hours * 3600 + minutes * 60 + seconds + milliseconds / 1000

def seconds_to_srt_time(seconds: float) -> str:
    ms = int((seconds % 1) * 1000)
    secs = int(seconds) % 60
    mins = int(seconds // 60) % 60
    hrs = int(seconds // 3600)
    return f"{hrs:02d}:{mins:02d}:{secs:02d},{ms:03d}"

def parse_srt(content: str):
    lines = [line.strip() for line in content.strip().splitlines() if line.strip() != '']
    segments = []
    i = 0
    while i < len(lines):
        if lines[i].isdigit():
            i += 1
            if i >= len(lines): break
            if '-->' in lines[i]:
                start_str, end_str = lines[i].split('-->')
                start = srt_time_to_seconds(start_str.strip())
                end = srt_time_to_seconds(end_str.strip())
                i += 1
                text_parts = []
                while i < len(lines) and not lines[i].isdigit():
                    text_parts.append(lines[i])
                    i += 1
                text = ' '.join(text_parts).strip()
                segments.append((start, end, text))
                continue
        i += 1
    return segments

def write_textgrid(segments, output_path, tier_name=None):
    if not segments:
        raise ValueError("无字幕数据")
    xmin = 0.0
    xmax = max(end for _, end, _ in segments)
    if tier_name is None:
        tier_name = "transcript"

    # 格式化时间值：如果为 0 则输出整数，否则保留三位小数
    def fmt_time(t):
        return '0' if t == 0.0 else f'{t:.3f}'

    with open(output_path, 'w', encoding='utf-8') as f:
        f.write('File type = "ooTextFile"\n')
        f.write('Object class = "TextGrid"\n\n')
        f.write(f'xmin = {fmt_time(xmin)}\n')
        f.write(f'xmax = {xmax:.3f}\n')
        f.write('tiers? <exists>\n')
        f.write('size = 1\n')
        f.write('item []:\n')
        f.write('    item [1]:\n')
        f.write('        class = "IntervalTier"\n')
        f.write(f'        name = "{tier_name}"\n')
        f.write(f'        xmin = {fmt_time(xmin)}\n')
        f.write(f'        xmax = {xmax:.3f}\n')
        f.write(f'        intervals: size = {len(segments)}\n')

        for idx, (start, end, text) in enumerate(segments, 1):
            f.write(f'        intervals [{idx}]:\n')
            f.write(f'            xmin = {fmt_time(start)}\n')
            f.write(f'            xmax = {end:.3f}\n')
            escaped_text = text.replace("\"", "\"\"")  # 转义双引号
            f.write(f'            text = "{escaped_text}"\n')

    print(f"✅ TextGrid 文件已生成：{output_path}")
    print(f"   总时长: {xmin:.3f} - {xmax:.3f} 秒")
    print(f"   字幕条数: {len(segments)}")

def parse_textgrid(content: str):
    lines = [line.strip() for line in content.splitlines() if line.strip()]
    segments = []
    in_intervals = False
    xmin = None
    xmax = None
    text = None

    for line in lines:
        if 'class = "IntervalTier"' in line:
            in_intervals = True
            continue
        if in_intervals and line.startswith('xmin ='):
            xmin = float(line.split('=', 1)[1].strip())
        elif in_intervals and line.startswith('xmax ='):
            xmax = float(line.split('=', 1)[1].strip())
        elif in_intervals and line.startswith('text ='):
            match = re.search(r'text = "(.*)"', line)
            text = match.group(1) if match else ""
            if xmin is not None and xmax is not None and text is not None:
                segments.append((xmin, xmax, text))
                xmin = xmax = text = None  # 重置用于下一个
    return segments

def write_srt(segments, output_path):
    with open(output_path, 'w', encoding='utf-8') as f:
        for idx, (start, end, text) in enumerate(segments, 1):
            start_str = seconds_to_srt_time(start)
            end_str = seconds_to_srt_time(end)
            f.write(f"{idx}\n")
            f.write(f"{start_str} --> {end_str}\n")
            f.write(f"{text}\n\n")

# ========== Gradio Web 界面 ==========

def srt_to_textgrid_file(srt_file, tier_name):
    try:
        with open(srt_file.name, 'r', encoding='utf-8') as f:
            content = f.read()
        segments = parse_srt(content)
        output_path = srt_file.name.rsplit('.', 1)[0] + '.TextGrid'
        write_textgrid(segments, output_path, tier_name or "transcript")
        return output_path
    except Exception as e:
        raise gr.Error(f"SRT 转换失败: {str(e)}")

def textgrid_to_srt_file(textgrid_file):
    try:
        with open(textgrid_file.name, 'r', encoding='utf-8') as f:
            content = f.read()
        segments = parse_textgrid(content)
        output_path = textgrid_file.name.rsplit('.', 1)[0] + '.srt'
        write_srt(segments, output_path)
        return output_path
    except Exception as e:
        raise gr.Error(f"TextGrid 转换失败: {str(e)}")

# 构建界面
with gr.Blocks(title="字幕转换器 | SRT ↔ TextGrid") as app:
    gr.Markdown("# 🔄 SRT 与 TextGrid 双向转换器")
    gr.Markdown("支持中文字幕，精准时间对齐，一键转换下载")

    with gr.Tabs():
        # Tab 1: SRT → TextGrid
        with gr.Tab("SRT → TextGrid"):
            with gr.Row():
                with gr.Column():
                    srt_upload = gr.File(label="上传 SRT 文件", file_types=['.srt'])
                    tier_name_input = gr.Textbox(label="TextGrid 层名（可选）", value="transcript")
                    srt_submit = gr.Button("转换为 TextGrid", variant="primary")
                with gr.Column():
                    srt_output = gr.File(label="下载 TextGrid 文件")
            srt_submit.click(
                fn=srt_to_textgrid_file,
                inputs=[srt_upload, tier_name_input],
                outputs=srt_output
            )

        # Tab 2: TextGrid → SRT
        with gr.Tab("TextGrid → SRT"):
            with gr.Row():
                with gr.Column():
                    tg_upload = gr.File(label="上传 TextGrid 文件", file_types=['.textgrid', '.TextGrid', '.txt'])
                    tg_submit = gr.Button("转换为 SRT", variant="primary")
                with gr.Column():
                    tg_output = gr.File(label="下载 SRT 文件")
            tg_submit.click(
                fn=textgrid_to_srt_file,
                inputs=tg_upload,
                outputs=tg_output
            )

    gr.Markdown("### 💡 使用说明")
    gr.Markdown("""
    - 支持拖拽上传 `.srt` 或 `.TextGrid` 文件
    - 转换后点击文件名即可下载
    - 时间精度保留毫秒级
    - 适用于 Praat 分析、字幕对齐等场景
    """)

# 启动服务器
if __name__ == "__main__":
    app.launch()